#!/usr/local/bin/python3
# *_* coding: UTF-8 *_*
# @IDE: PyCharm
# @Version: Python3.7
# @Author: Kendrick.Kc
# @Email: 509556985@qq.com
# @File: ERA.py
# @Inst: ERA策略(忽略15分钟的周期)
# @Time: 2023/2/3 14:55
# -----


import backtest
import pandas
pandas.set_option('expand_frame_repr', False)  # 当列太多时不换行
pandas.set_option('display.max_rows', 10000)  # 最多显示数据的行数


class ERA(backtest.VBT):
    # start_time = "2023-01-25"  # 回测数据的开始日期
    # end_time = "2023-08-25"  # 回测数据的结束日期
    kcEma = 72  # EMA参数. 目前最优:72
    kcAtr = 20  # ATR参数. 目前最优:20
    kcRsi = 12  # RSI参数. 目前最优:12
    kcRsi7 = 72  # RSI参数. 目前最优:72
    kcRsi2 = 22  # RSI参数. 目前最优:22
    kcAtrMultiple = 0.1     # ATR倍数. 目前最优:0.1
    # ---------------------- 上面是v1

    # 用于记录指标：指标未记录时，不交易
    rsi = None
    shang_gui = None
    xia_gui = None
    date_15 = None

    def clean_data(self):
        """
        计算指标
        :return:
        """
        data_15 = self.get_data(bar="15m")
        data_15["ema"] = backtest.pandas_ta.ema(data_15["close"], self.kcEma)
        data_15["atr"] = backtest.pandas_ta.atr(data_15["high"], data_15["low"], data_15["close"], self.kcAtr)
        data_15["rsi"] = backtest.pandas_ta.rsi(data_15["close"], self.kcRsi)
        data_15["shang_gui"] = data_15["ema"] + self.kcAtrMultiple * data_15["atr"]
        data_15["xia_gui"] = data_15["ema"] - self.kcAtrMultiple * data_15["atr"]

        self.data["rsi"] = data_15["rsi"]
        self.data["shang_gui"] = data_15["shang_gui"]
        self.data["xia_gui"] = data_15["xia_gui"]
        self.data["date_15"] = data_15["date"]
        # ---------------------- 上面是v1

        # title = ['close', 'rsi', 'shang_gui', 'xia_gui']
        # print(self.data[title].tail(100))
        # print(data_15[title].tail(100))
        # exit()

    def strategy(self, data):
        """
        策略
        :param data: 当前1分钟Kbar数据
        :return:
        """

        # 如果为当前15分钟，那么就保存上下轨并跳出当前分钟从下一分钟开始计算
        if not backtest.numpy.isnan(data.rsi):
            self.shang_gui, self.xia_gui, self.rsi = data.shang_gui, data.xia_gui, data.rsi
            self.date_15 = data.date_15
            return

        print(f"############## 1分钟时间={data.date}, 15分钟时间={self.date_15}")

        if not self.long_pos and not self.short_pos and self.rsi is not None:

            if data.close > self.shang_gui and self.rsi > self.kcRsi7:
                self.entries_pos(self.buy_long)
                self.long_pos_exits_price = self.xia_gui

            if data.close < self.xia_gui and self.rsi < self.kcRsi2:
                self.entries_pos(self.sell_short)
                self.short_pos_exits_price = self.shang_gui

        if self.long_pos:
            if data.close < self.long_pos_exits_price:
                self.exits_pos(self.sell_long)
            else:
                self.long_pos_exits_price = self.xia_gui

        if self.short_pos:
            if data.close > self.short_pos_exits_price:
                self.exits_pos(self.buy_short)
            else:
                self.short_pos_exits_price = self.shang_gui

    def finish(self):

        # 将stats结果写入文件
        file = open(self.current_path_splice(f"{str(int(backtest.time.time() * 1000))}.txt"), 'w')
        file.write(str(self.pf.stats()))
        file.close()


run = ERA()
run.run_strategy(run.strategy)
